#========================================
#    author: Changlong.Zang
#      mail: zclongpop123@163.com
#      time: Tue Apr 18 10:49:47 2017
#========================================
import maya.cmds as mc
import stUtils
#--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
def get_reference_namespace():
    '''
    '''
    for ref in mc.file(q=True, r=True):
        try:
            yield mc.file(ref, q=True, ns=True)
        except:
            pass





def get_refrence_nodes():
    '''
    '''
    data = dict()

    for rn in mc.ls(typ='reference'):
        if rn == 'sharedReferenceNode':
            continue
        try:
            namesps = mc.referenceQuery(rn, ns=True)[1:]
            data[namesps] = rn
        except:
            continue

    return data





def get_reference_controls(namespace):
    '''
    '''
    ref_node_dict = get_refrence_nodes()
    ref_rn_node   = ref_node_dict.get(namespace)
    if not ref_rn_node:
        return list()

    reference_nodes = mc.referenceQuery(ref_rn_node, dp=True, n=True)
    if not reference_nodes:
        return list()

    reference_controls = mc.listRelatives(mc.ls(reference_nodes, typ='nurbsCurve'), p=True, path=True)
    reference_controls = [x for i, x in enumerate(reference_controls) if x not in reference_controls[:i]]

    return reference_controls





def get_reference_keys():
    '''
    '''
    data = dict()

    ref_namespace_list = get_reference_namespace()

    for namespace in ref_namespace_list:
        ref_main_ctl = '{0}:Main'.format(namespace)
        if not mc.objExists(ref_main_ctl):
            continue

        if not mc.attributeQuery('OrigAnimTag', ex=True, n=ref_main_ctl):
            continue

        orig_anim_keys = mc.keyframe('{0}.OrigAnimTag'.format(ref_main_ctl), q=True) or list()
        data['{0}_Orig'.format(namespace)] = [f for f in orig_anim_keys if mc.getAttr('{0}.OrigAnimTag'.format(ref_main_ctl), t=f) == 0]
        data['{0}_Anim'.format(namespace)] = [f for f in orig_anim_keys if mc.getAttr('{0}.OrigAnimTag'.format(ref_main_ctl), t=f) == 1]

    return data





@stUtils.undo_decorator
def key_to_step(namespace, step=1, start=None, end=None):
    '''
    '''
    ns = namespace
    RN = get_refrence_nodes()[ns]

    #-
    names  = mc.keyframe(mc.referenceQuery(RN, n=True, dp=True), q=True, n=True) or list()
    names  = [n for n in names if not mc.referenceQuery(n, inr=True)]
    if not names:
        return

    #-
    frames = mc.keyframe(names, q=True, tc=True) or list()
    frames = [int(x) for i, x in enumerate(frames) if x not in frames[:i]]
    frames.sort()
    if len(frames) < 2:
        return

    #-
    if start is None and end is None:
        mc.keyframe(names, time=(frames[0], frames[-1]), r=True, tc=10000, absolute=True)
        for i, f in enumerate(frames):
            new_time = min(frames) + i * step
            mc.keyframe(names, time=(f + 10000,), tc=new_time, absolute=True)

    else:
        anim_start = mc.playbackOptions(q=True, ast=True)
        mc.keyframe(names, time=(anim_start + start, frames[-1]), r=True, tc=10000, absolute=True)
        i = 0
        for f in frames:
            if f < (anim_start + start):
                continue

            if f > (anim_start + end):
                continue

            new_time = anim_start + start + i * step
            mc.keyframe(names, time=(f + 10000,), tc=new_time, absolute=True)

            i += 1

        mc.keyframe(names, time=(anim_start + end + 10000, frames[-1]+10000), tc=new_time + 1 - (anim_start + end + 10000), r=True, absolute=True)





@stUtils.undo_decorator
def paste_keys(src, dst):
    '''
    '''
    src_objects = get_reference_controls(src[0].rsplit('_', 1)[0])
    if not src_objects:
        return

    start   = mc.playbackOptions(q=True, ast=True)
    if src[2] == 'Copy':
        mc.copyKey(src_objects, time=(start + min(src[1]), start + max(src[1])))

    elif src[2] == 'Cut':
        mc.cutKey(src_objects, time=(start + min(src[1]), start + max(src[1])))

    else:
        pass

    dst_objects = get_reference_controls(dst[0].rsplit('_', 1)[0])
    if not dst_objects:
        return

    mc.pasteKey(dst_objects, t=(start + dst[1][0], start + dst[1][0]))





@stUtils.undo_decorator
def delete_keys(namespace, keys):
    '''
    '''
    objects = get_reference_controls(namespace.rsplit('_', 1)[0])
    start   = mc.playbackOptions(q=True, ast=True)
    for k in keys:
        mc.cutKey(objects, t=(start+k, start+k), cl=True)





def set_orig_anim_key(namespace, frame, typ):
    '''
    '''
    ref_main_ctl = '{0}:Main'.format(namespace)
    if mc.objExists(ref_main_ctl) and mc.attributeQuery('OrigAnimTag', ex=True, n=ref_main_ctl):
        if typ == 'Orig':
            mc.setKeyframe('{0}.OrigAnimTag'.format(ref_main_ctl), t=frame, v=0)
        else:
            mc.setKeyframe('{0}.OrigAnimTag'.format(ref_main_ctl), t=frame, v=1)





@stUtils.undo_decorator
def set_keys(namespace, frame, typ='Orig'):
    '''
    '''
    controls = get_reference_controls(namespace)
    if not controls:
        return
    mc.setKeyframe(controls, t=frame)

    set_orig_anim_key(namespace, frame, typ)





@stUtils.undo_decorator
def move_keys(src_namespace, src_frame, dst_namespace, dst_frame):
    '''
    '''
    src_ns = src_namespace.rsplit('_', 1)[0]
    dst_ns = dst_namespace.rsplit('_', 1)[0]

    if src_ns == dst_ns and src_frame != dst_frame:
        src_controls = get_reference_controls(src_ns)
        mc.keyframe(src_controls, time=(src_frame,), tc=dst_frame, a=True)

    typ = dst_namespace.rsplit('_', 1)[-1]
    set_orig_anim_key(src_ns, dst_frame, typ)





def bake_char_anim(namespace):
    '''
    '''
    objects = get_reference_controls(namespace)
    start = mc.playbackOptions(q=True, ast=True)
    end   = mc.playbackOptions(q=True, aet=True)
    mc.bakeResults(objects, t=(start, end), simulation=True)
